package edu.cn.nlsde.tmfst.lda;

import edu.cn.nlsde.Utils.IOUtil;

import java.io.BufferedReader;
import java.io.IOException;
import java.util.*;
import java.util.Map.Entry;

public class TmpTopicModel {

    public Map<Integer, String> id2word = new HashMap<Integer, String>();
    public Map<String, Integer> word2id = new HashMap<String, Integer>();

    public int[][] docs;

    public int[][] zAssigns;

    public int[][] ndk;
    public int[] ndkSum;

    public int[][] nkw;
    public int[] nkwSum;
    public double[] pk;
    public double alpha = 0.1;
    public double beta = 0.01;
    public double u = .05;
    public double delta = .1;
    public int M;
    public int V;
    public int K = 200;
    public int KG = 10;
    Map<Integer, Set<Integer>> k2dSet = new HashMap<Integer, Set<Integer>>();

    public static void main(String args[]) {
        TmpTopicModel ttm = new TmpTopicModel();
        ttm.loadData("/Users/zuoyuan/Desktop/experiment/20130601-train");
        ttm.initModel();
        ttm.doSampling(200);
        ttm.print_topics(10);
    }

    public void loadData(String path) {
        List<String> docsBuff = new ArrayList<String>();

        BufferedReader reader = IOUtil.getReader(path, "utf-8");

        String line;
        try {
            line = reader.readLine();
            while (line != null) {
                docsBuff.add(line.trim());
                line = reader.readLine();
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        M = docsBuff.size();
        docs = new int[M][];

        for (int m = 0; m != M; m++) {
            String[] tokens = docsBuff.get(m).trim().split("\\s+");

            int N = tokens.length;
            docs[m] = new int[N];

            for (int n = 0; n != N; n++) {
                if (!word2id.containsKey(tokens[n])) {
                    word2id.put(tokens[n], word2id.size());
                    id2word.put(word2id.get(tokens[n]), tokens[n]);
                }
                docs[m][n] = word2id.get(tokens[n]);
            }
        }

        V = word2id.size();

        return;
    }

    public void initModel() {
        zAssigns = new int[M][];

        ndk = new int[M][K];
        ndkSum = new int[M];

        nkw = new int[K][V];
        nkwSum = new int[K];

        Random rand = new Random(System.currentTimeMillis());

        for (int m = 0; m != M; m++) {
            int N = docs[m].length;
            zAssigns[m] = new int[N];

            for (int n = 0; n != N; n++) {
                int w = docs[m][n];
                int z = rand.nextInt(K);

                Set<Integer> docSet;
                if (!k2dSet.containsKey(z))
                    docSet = new HashSet<Integer>();
                else
                    docSet = k2dSet.get(z);
                docSet.add(m);

                k2dSet.put(z, docSet);

                ndk[m][z]++;
                ndkSum[m]++;

                nkw[z][w]++;
                nkwSum[z]++;

                zAssigns[m][n] = z;
            }
        }

        pk = new double[K];
        for (int k = 0; k != K; k++) {
            updatePK(k);
        }
    }

    public void doSampling(int niters) {
        for (int iter = 0; iter != niters; iter++) {
            System.out.println("Iteration " + iter);
            for (int m = 0; m != M; m++) {
                for (int n = 0; n != docs[m].length; n++) {
                    sample(m, n);
                }
            }
        }
    }

    public void sample(int m, int n) {
        int z = zAssigns[m][n];
        int w = docs[m][n];

        ndk[m][z]--;
        ndkSum[m]--;

        nkw[z][w]--;
        nkwSum[z]--;

        Set<Integer> docSet;
        if (ndk[m][z] == 0) {
            docSet = k2dSet.get(z);
            docSet.remove(m);
            k2dSet.put(z, docSet);
            updatePK(z);
        }

        double[] pTable = new double[K];
        double pTotal = 0.0;

        for (int k = 0; k != K; k++) {
            if (k < KG) {
                pTable[k] = (ndk[m][k] + alpha) * (nkw[k][w] + beta) / (nkwSum[k] + V * beta);
            } else {
                pTable[k] = pk[k] * (ndk[m][k] + alpha) * (nkw[k][w] + beta) / (nkwSum[k] + V * beta);
            }
            pTotal += pTable[k];
        }

        double u = Math.random() * pTotal;
        double v = 0.0;

        for (int k = 0; k != K; k++) {
            v += pTable[k];
            if (v > u) {
                z = k;
                break;
            }
        }

        ndk[m][z]++;
        ndkSum[m]++;

        nkw[z][w]++;
        nkwSum[z]++;

        if (ndk[m][z] == 1) {
            docSet = k2dSet.get(z);
            docSet.add(m);
            k2dSet.put(z, docSet);
            updatePK(z);
        }

        zAssigns[m][n] = z;
    }

    private void updatePK(int k) {
        double nkd = k2dSet.get(k).size();
        pk[k] = Math.exp((-(nkd / M - u) * (nkd / M - u) / (2 * delta * delta)));
    }

    public double[][] compute_phi() {
        double[][] phi = new double[K][V];
        for (int k = 0; k != K; k++)
            for (int w = 0; w != V; w++)
                phi[k][w] = (nkw[k][w] + beta) / (nkwSum[k] + V * beta);
        return phi;
    }

    public ArrayList<List<Entry<String, Double>>> sorted_topicwords(
            double[][] phi, int T) {
        ArrayList<List<Entry<String, Double>>> res = new ArrayList<List<Entry<String, Double>>>();
        for (int k = 0; k != T; k++) {
            HashMap<String, Double> term2weight = new HashMap<String, Double>();
            for (String term : word2id.keySet())
                term2weight.put(term, phi[k][word2id.get(term)]);

            List<Entry<String, Double>> pairs = new ArrayList<Entry<String, Double>>(
                    term2weight.entrySet());
            Collections.sort(pairs, new Comparator<Entry<String, Double>>() {
                public int compare(Entry<String, Double> o1,
                                   Entry<String, Double> o2) {
                    return (o2.getValue().compareTo(o1.getValue()));
                }
            });

            res.add(pairs);
        }

        return res;
    }

    public void print_topics(int top_n) {
        double phi[][] = this.compute_phi();
        ArrayList<List<Entry<String, Double>>> pairs_list = this
                .sorted_topicwords(phi, K);
        for (int k = 0; k != K; k++) {
            System.out.println("Topic " + k + ":");
            for (int i = 0; i != top_n; i++) {
                System.out.println(pairs_list.get(k).get(i).getKey() + " "
                        + pairs_list.get(k).get(i).getValue());
            }
        }
    }
}
